#include <iostream>
#include <opencv2/opencv.hpp>
#include <io.h>  // 遍历文件夹文件
#include <opencv2/core/utils/logger.hpp>

#include "trt_run.h"

YoloDetect detect_person(const std::string &engine_file, std::vector<std::string> &class_names) {
    // 检测类别
    class_names = {"person"};
    // 一个关键点几个值(2个[x,y],3个[x,y,v]，v是是否可见)
    const int point_dim = 3;
    const float conf_threshold = 0.3f;
    // 是否画框
    const bool draw_box = true;

    // 关键点连线关系
    const std::vector<std::array<int, 2>> drawLine = {{0,  1},
                                                      {0,  2},
                                                      {1,  3},
                                                      {2,  4},
                                                      {3,  5},
                                                      {4,  6},
                                                      {5,  7},
                                                      {6,  8},
                                                      {7,  9},
                                                      {8,  10},
                                                      {5,  6},
                                                      {5,  11},
                                                      {6,  12},
                                                      {11, 12},
                                                      {11, 13},
                                                      {12, 14},
                                                      {13, 15},
                                                      {14, 16}};

    // 模板颜色
    std::vector<cv::Scalar> colorTemp = {{255, 128, 0},
                                                 {255, 51,  255},
                                                 {51,  153, 255},
                                                 {0,   255, 0}};

    // 关键点颜色
    std::vector<int> kptIndex = {3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2};
    std::vector<cv::Scalar> kptColor;
    kptColor.reserve(kptIndex.size());
    for (int i: kptIndex) {
        kptColor.push_back(colorTemp[i]);
    }

    // 连线颜色
    std::vector<int> lineIndex = {3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2};
    std::vector<cv::Scalar> lineColor;
    lineColor.reserve(lineIndex.size());
    for (int i: lineIndex) {
        lineColor.push_back(colorTemp[i]);
    }

    // 使用 new 关键字在堆上分配内存
    auto *yolo_detect = new YoloDetect(engine_file, conf_threshold, class_names, point_dim, drawLine, kptColor,
                                       lineColor, draw_box);
    // 返回对象的引用
    return *yolo_detect;
};

int main() {
    //LOG_LEVEL_ERROR - 只输出错误日志 / LOG_LEVEL_SILENT - 不在输出日志
    cv::utils::logging::setLogLevel(cv::utils::logging::LOG_LEVEL_ERROR);

    // 图像全路径单斜杠需要转义，或者需要R"()"这样
    // 检测路径
    std::string images_path = R"(E:\new_code\github\dataset\coco128-pose\images\train2017)";
    // 模型路径
    std::string engine_file = R"(G:\Github\yolov8_20231114\ultralytics\my_test\weights\yolov8_pose\yolov8s-pose_int8.engine)";

    // 存储文件信息的结构体
    struct _finddata_t file_info{};
    std::string current_path = images_path + "/*.jpg";  // 查找特定后缀的文件，支持通配符
    // 第一次查找，查找成功返回longlong型的句柄，失败返回-1
    long long handle = _findfirst(current_path.c_str(), &file_info);

    // 类别
    std::vector<std::string> class_names;
    // 调用构造函数，初始化模型
    YoloDetect yolo_detect = detect_person(engine_file, class_names);

    // 存储画框图
    cv::Mat img;
    cv::Mat draw_img;
    // 存储输出结果
    std::vector<Detection> result;

    double start, end, t;
    do {
//        // 单张图像路径
        std::string img_path = images_path + "/" + file_info.name;
//
        // Mat是创建一个数组
        // IMREAD_COLOR/1-默认彩色  IMREAD_GRAYSCALE/0-灰度显示
        // IMREAD_UNCHANGED/-1-四通道png
        img = cv::imread(img_path, 1);
        draw_img = img;
//        std::cout << draw_img.type() << std::endl;
//        std::cout<<draw_img.size<<std::endl;  // 2048 x 2448
//        std::cout<<draw_img.channels()<<std::endl;  // 3
//        std::cout<<draw_img.dims<<std::endl;  // 2

        start = clock();
//        for (int i = 0; i < 100; i++) {
//            // 侦测
//            result = yolo_detect.detect(draw_img);
//        }
        // 侦测
        result = yolo_detect.detect(draw_img);

        end = clock();
        // 计算毫秒ms
        t = (end - start) / CLOCKS_PER_SEC * 1000;
        std::cout << "detect_time: " << t << std::endl;

        std::cout << "out_boxes: " << result.size() << std::endl;
//        for (const Detection &detection: result) {
//            int clas_index = detection.clas_id;
//            const std::string &label = class_names[clas_index];
//            float conf = (float) std::round(detection.conf * 100) / 100;
//            int x = detection.box.x;
//            int y = detection.box.y;
//            int w = detection.box.width;
//            int h = detection.box.height;
//
//            std::cout << label << ": " << conf << " box(x,y,w,h): " << x << ", " << y << ", " << w << ", " << h
//                      << std::endl;
//
//            std::vector<std::tuple<int, int, float>> keypoints = detection.keypoints;
//            for (const auto &keypoint: keypoints) {
//                int px = std::get<0>(keypoint);
//                int py = std::get<1>(keypoint);
//                float pconf = std::get<2>(keypoint);
//                std::cout << "keypoint(x,y,conf): " << px << ", " << py << ", " << pconf << std::endl;
//            }
//
//        }

        std::cout << "============================" << std::endl;

//        cv::namedWindow("show_img", cv::WINDOW_NORMAL);
//        cv::imshow("show_img", draw_img);
//        cv::waitKey(0);

    } while (!_findnext(handle, &file_info));  // 查找下一个
    _findclose(handle);  // 关闭

    delete &yolo_detect;
    cv::destroyAllWindows();  // 销毁窗口
    system("pause");
}